from torchvision import models
import torch.nn as nn
import torch
import torch.nn.functional as F
from collections import defaultdict


class HiLo_Transform(torch.nn.Module):
    def __init__(self, M, N) -> None:
        super().__init__()
           # 提取高频成分
        crow, ccol = M // 2, N // 2
        self.pos = torch.ones((M, N, 2)).cuda()
        for x in range(M):
            for y in range(N):
                self.pos[x, y, 0] = -(x-crow)**2
                self.pos[x, y, 1] = -(y-ccol)**2
        
        self.pos = self.pos.reshape(-1, 2)
        self.pos.requires_grad = False


    def forward(self, img):
        b, _, M, N = img.shape
        sigma = 0.5*torch.ones(b, 3, 2).cuda()
        # sigma (batchsize, 3, 2) self.mask_high_pass (M×N, 2)
        sigma = sigma.reshape(-1, 2) # sigma (batchsize×3, 2)
        
        low_pass = sigma @ self.pos.T
        low_pass = low_pass.reshape(b, 3, M, N)
        low_pass = low_pass.exp()
        high_pass = 1-low_pass

        freq_ = torch.fft.fft2(img)
        freq = torch.fft.fftshift(freq_)
        high = high_pass * freq  
        low = low_pass * freq

        high = torch.fft.ifftshift(high)   
        high = torch.fft.ifft2(high)

        low = torch.fft.ifftshift(low)   
        low = torch.fft.ifft2(low)
        # return ori_img.real
        return torch.abs(high), torch.abs(low)

def idx_init(a, b, c):
    def cal_idx(i, j, k):
        return i*(b*c) + j*c + k
    return cal_idx

class hilo(nn.Module):
    def __init__(self, num_class=2):
        super().__init__()
        self.in_channels = 3
        self.hilo_transform = HiLo_Transform(224, 224)
        cfg = [[256, "M"], [256, "M"], [256, 256, "M"], [256, 256, "M"], [256, 256, "M"]]
        self.encoder = nn.ModuleList([])
        self.multi_scale_fusion = nn.ModuleList([])
        for c in cfg:
            self.encoder.append(self._make_layers(c))
        self.avgpool = nn.AdaptiveAvgPool3d((512, 7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Linear(1024, num_class)
        )
        self.graph = defaultdict(lambda: [])
        edge = [nn.Identity()]*(4*4*4)
        cal_idx = idx_init(4, 4, 4)
        for i in range(1, 4):
            for j in range(4):
                for k in range(4):
                    self.graph[i, j].append((i-1, k))
                    if j == k:
                        edge[cal_idx(i-1,k,j)] = nn.Conv2d(256, 256, 3, 1, 1)
                    elif j < k:
                        # 上采样
                        edge[cal_idx(i-1,k,j)] = torch.nn.Upsample(scale_factor=2**(k-j), mode='bilinear')
                    else:
                        # 下采样
                        edge[cal_idx(i-1,k,j)] = nn.MaxPool2d(2**(j-k), 2**(j-k))
        i = 4
        j = 0
        for k in range(4):
            self.graph[i, j].append((3, k))
            if j == k:
                edge[cal_idx(i-1,k,j)] = nn.Conv2d(256, 256, 3, 1, 1)
            elif j < k:
                # 上采样
                edge[cal_idx(i-1,k,j)] = torch.nn.Upsample(scale_factor=2**(k-j), mode='bilinear')
            else:
                # 下采样
                edge[cal_idx(i-1,k,j)] = nn.MaxPool2d(2**(j-k), 2**(j-k))
            
        self.edge = nn.ModuleList(edge)
            
        
    def forward(self, x):
        cal_idx = idx_init(4, 4, 4)
        def dfs(node, lst):
            i, j = node
            if i == 0:
                return lst[j]
            tmp = []
            for pre_node in self.graph[node]:
                out = self.edge[cal_idx(pre_node[0],pre_node[1],node[1])](dfs(pre_node, lst))
                tmp.append(out)
            return sum(tmp)

        lo_lst = []
        hi_lst = []
        lo_freq_features, hi_freq_features = self.hilo_transform(x)
        # 224
        lo_lst = []
        hi_lst = []
        for i, layer in enumerate(self.encoder):
            lo_freq_features = layer(lo_freq_features)
            hi_freq_features = layer(hi_freq_features)
            if i != 0:
                lo_lst.append(lo_freq_features)
                hi_lst.append(hi_freq_features)

     
        lo_freq_features = dfs((4, 0), lo_lst)

        hi_freq_features = dfs((4, 0), hi_lst)

        
        x = torch.cat([lo_freq_features,  hi_freq_features], dim=1)
       
        x = self.avgpool(x)
        # 改变tensor形状，拉伸成一维
        x = x.view(x.size(0), -1)
        out = self.classifier(x)
        return out
    
    def _make_layers(self, cfg, batch_norm: bool = False) -> nn.Sequential:
        layers = []
        # self.in_channels = 3
        for v in cfg:
            if v == "M":
                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
            else:
                layers.append(nn.Conv2d(self.in_channels, v, kernel_size=3, padding=1))
                if batch_norm:
                    layers.append(nn.BatchNorm2d(v))
                layers.append(nn.ReLU(inplace=True))
                self.in_channels = v
        return nn.Sequential(*layers)
    
        
if __name__ == '__main__':
    model = hilo().cuda()
    x = torch.randn(10, 3, 224, 224).cuda()
    y = model(x)
    print(y.shape)
